package com.leige.datasource.core;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import javax.sql.DataSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.JndiDataSourceLookup;
import org.springframework.util.Assert;

/**
 * 配置主从数据源后，根据选择，返回对应的数据源。多个从库的情况下，会平均的分配从库，用于负载均衡。
 *
 * @author tanghd
 */

public class DynamicDataSource extends AbstractDataSource implements ApplicationContextAware {

    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);
	

    private static List<String> slavesDataSourceNames;
    private static String defaultTargetJndiName;
    private static final ThreadLocal<LinkedList<Boolean>> ismaster=new ThreadLocal<LinkedList<Boolean>>(){protected LinkedList<Boolean> initialValue(){ return new LinkedList<Boolean>();}};
    
    private static int selectnum=0;
    
    public static void useMaster() { 
    	ismaster.get().push(true); 
    }

    public static void useSlave() { 
    	ismaster.get().push(false); 
    }

    public static void reset() { 
    	ismaster.get().pop();
    	if(ismaster.get().size()==0)
    		ismaster.remove();
    }
    
	private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
	private Map<String, DataSource> resolvedDataSources=new LinkedHashMap<String,DataSource>();;
	private DataSource resolvedDefaultDataSource;
 
	public void setDefaultTargetDataSource(DataSource defaultTargetDataSource) {
		this.resolvedDefaultDataSource = defaultTargetDataSource;
	}
	public void setDefaultTargetJndiName(String defaultTargetJndiName){
		DynamicDataSource.defaultTargetJndiName=defaultTargetJndiName;
	}
	public void setDataSourceLookup(DataSourceLookup dataSourceLookup) {
		this.dataSourceLookup = (dataSourceLookup != null ? dataSourceLookup : new JndiDataSourceLookup());
	}
	
	public void setSlavesDataSourceNames(List<String> slavesDataSourceNames) {
		DynamicDataSource.slavesDataSourceNames=slavesDataSourceNames;
	}

	@Override
	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
		// TODO Auto-generated method stub
		if(resolvedDefaultDataSource==null&&defaultTargetJndiName!=null)
			try {
				resolvedDefaultDataSource=this.dataSourceLookup.getDataSource(defaultTargetJndiName);
			} catch (Exception e) {
				logger.warn("查找DataSource 从jndi中查找"+defaultTargetJndiName+"未找到",e);
			}
		if(slavesDataSourceNames==null) return;
		for (String slavesDataSourceName : slavesDataSourceNames) {
			try {
				resolvedDataSources.put(slavesDataSourceName, applicationContext.getBean(slavesDataSourceName,DataSource.class));
			} catch (BeansException e) {
				logger.warn("查找DataSource 从spring上下文中查找"+slavesDataSourceName+"未找到",e);
				try {
					resolvedDataSources.put(slavesDataSourceName, this.dataSourceLookup.getDataSource(slavesDataSourceName));
				} catch (Exception e2) {
					logger.warn("查找DataSource 从jndi中查找"+slavesDataSourceName+"未找到",e);
				}
			}
		}
	}

	@Override
	public Connection getConnection() throws SQLException {
		return determineTargetDataSource().getConnection();
	}

	@Override
	public Connection getConnection(String username, String password) throws SQLException {
		return determineTargetDataSource().getConnection(username, password);
	}

	@Override
	@SuppressWarnings("unchecked")
	public <T> T unwrap(Class<T> iface) throws SQLException {
		if (iface.isInstance(this)) {
			return (T) this;
		}
		return determineTargetDataSource().unwrap(iface);
	}

	@Override
	public boolean isWrapperFor(Class<?> iface) throws SQLException {
		return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
	}

	protected DataSource determineTargetDataSource() {
		Object lookupKey = determineCurrentLookupKey();
		DataSource dataSource = this.resolvedDataSources.get(lookupKey);
		if (dataSource != null)
			return dataSource;
		return this.resolvedDefaultDataSource;
	}

	
    /**
     * 如果是选择使用从库，且从库的数量大于1，则通过取模来控制从库的负载,
     */
    protected Object determineCurrentLookupKey() {
    	if(ismaster.get().size()==0||ismaster.get().getFirst())
    		return null;
    	if(slavesDataSourceNames!=null){
    		String slavesDataSourceName = slavesDataSourceNames.get(selectnum=(++selectnum %(slavesDataSourceNames.size())));
    		if(logger.isDebugEnabled())logger.debug("切换到从库"+slavesDataSourceName+"中查询");
    		return slavesDataSourceName;
    	}
    	return null;
    }

}
